package com.peiqi.aiproject.service;

import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.druid.sql.visitor.functions.Char;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.net.URI;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;

//继承TextWebSocketHandler文本发送处理器
@Component
public class ChatHandler extends TextWebSocketHandler {
    @Value("${ai.api-key}")
    private String appKey;
    @Resource
    private Generation generation;

    private static final Logger logger = LoggerFactory.getLogger(ChatHandler.class);

    private final Map<String, WebSocketSession> userIdToSessionMap = new ConcurrentHashMap<>();


//    前端链接成功时获取和重写方法
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = extractUserIdFromSession(session); // 实现此方法以从WebSocketSession中提取用户ID
        logger.info("New WebSocket connection established for user {}: {}", userId, session.getId());
        userIdToSessionMap.put(userId, session);
    }

    //    前端链接结束时和重写方法
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String userId = extractUserIdFromSession(session);
        logger.info("WebSocket connection closed for user {}: {}, Status: {}", userId, session.getId(), status.getCode());
        userIdToSessionMap.remove(userId, session);
    }

    /**
     * 发送消息给指定用户。
     *
     * @param recipientUserId 目标用户ID
     * @param message         要发送的消息内容
     */
//    发送 AI 流式输出
    public void sendMessageToUser(String recipientUserId, String message) {
        WebSocketSession targetSession = userIdToSessionMap.get(recipientUserId);
        if (targetSession != null) {
            try {
//                发送 AI 生成的 数据到客户端
                targetSession.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                logger.error("Failed to send message to user {}: {}", recipientUserId, e.getMessage());
            }
        } else {
            logger.warn("No active WebSocket session found for user {}", recipientUserId);
        }
    }

    /**
     * 从WebSocketSession中提取用户ID。实现此方法以适应您的具体应用逻辑。
     *
     * @param session WebSocketSession对象
     * @return 用户ID
     */
    private String extractUserIdFromSession(WebSocketSession session) {
        URI uri = session.getUri();
        String query = uri.getQuery();

        if (query != null) {
            String[] params = query.split("&");
            for (String param : params) {
                String[] parts = param.split("=");
                if (parts.length == 2 && "userId".equals(parts[0])) {
                    return parts[1];
                }
            }
        }

        throw new IllegalArgumentException("Missing or invalid 'userId' parameter in WebSocket connection URL.");
    }


    /**
     * 处理客户端发送的文本消息。
     *
     * @param session WebSocket会话
     * @param message 客户端发送的文本消息
     * @throws Exception 在处理消息过程中发生的任何异常
     */
//    处理客户端发送过来的文本
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String userId = extractUserIdFromSession(session);
//        获取文本信息
        String receivedMessage = message.getPayload();
        logger.info("Received message from user {}: {}", userId, receivedMessage);

        // 在此处根据业务需求处理收到的文本消息，例如：
        // 1. 将消息转发给其他在线用户
        // 2. 存储消息到数据库
        // 3. 调用AI服务进行回复
//        调用 通义千问 大模型的 API
        sendMessage(receivedMessage);
    }
    public void sendMessage(String question) throws NoApiKeyException {
        if(StrUtil.isBlank(appKey)){
            throw new NoApiKeyException();
        }
        // 构建用户消息
        Message message = Message.builder()
//                设置发送对象为用户，枚举类型标识的是用户
                .role(Role.USER.getValue())
                .content(question).build();

        // 设置AI对话参数
        QwenParam qwenParam = QwenParam.builder()
//                设置AI模型为 通义千问-Plus
                .model(Generation.Models.QWEN_PLUS)
//                添加消息列表（用户传进来的提问数据）
                .messages(Collections.singletonList(message))
//                设置ai的采样策略 (个人理解为 AI 的智商)
                .topP(0.8)
//                设置返回值为枚举类型消息 (MESSAGE)
                .resultFormat(QwenParam.ResultFormat.MESSAGE)
//                启用搜索功能, 使得AI回复附带参考文献 (比如学java推荐什么什么书)
                .enableSearch(true)
//                传入之前设置好的 API-KEY
                .apiKey(appKey)
//                启用增量输出，如果要回的内容长，就逐步返回结果,而不是全部输出再返回，避免程序阻塞
                .incrementalOutput(true)
                .build();
        Semaphore semaphore = new Semaphore(0);
        StringBuilder fullContent = new StringBuilder();
        // 调用AI生成服务
        CompletableFuture.runAsync(() -> {
            try {
                //        大模型的回调函数
                generation.streamCall(qwenParam, new ResultCallback<GenerationResult>() {

//                    监听回调事件
                    @Override
                    public void onEvent(GenerationResult message) {
                        String content = message.getOutput().getChoices().get(0).getMessage().getContent();
//                        使用 空字符串 替换 *
                        String result = content.replace("*", "");
                        fullContent.append(result);
//                        不断拼接 ai 发送的数据信息
                        sendMessageToUser("1",fullContent.toString());
                        System.out.println(message);
                    }
//                    监听回调错误
                    @Override
                    public void onError(Exception err) {
                        System.out.println(String.format("Exception: %s", err.getMessage()));
                        semaphore.release();
                    }
//                    监听回调数据
                    @Override
                    public void onComplete() {
                        System.out.println("Completed");
//                        发送 AI生成的 数据 到客户端
                        sendMessageToUser("1","End");
                        semaphore.release();
                    }

                });
            } catch (NoApiKeyException e) {
                throw new RuntimeException(e);
            } catch (InputRequiredException e) {
                throw new RuntimeException(e);
            }
        }, Executors.newSingleThreadExecutor());
//        semaphore.acquire();
    }
}
